#!/usr/bin/env python3
# coding: utf-8
# Copyright 2020 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ===========================================================================

import os
import platform
import shutil
import subprocess
from typing import List, Optional

from . import logger_config
from .download_data import DownloadData
from .download_util import get_obs_downloader_path,DownloadError
from .parallel_file_downloader import DownloadFileInfo
from .software_mgr import PkgInfo, SoftwareVersion
from ascend_deployer.module_utils.safe_file_handler import SafeFileHandler


LOG = logger_config.LOG
CUR_DIR = os.path.dirname(os.path.realpath(__file__))
PROJECT_DIR = os.path.dirname(CUR_DIR)


class DockerImagesDownloader:
    _IMAGES = ("MIS-image","GPUStack","vLLM-Ascend","openGaussImage","DifyImages")
    
    def __init__(self, download_data: DownloadData):
        self._download_data = download_data
        self._arch = download_data.arch
        self._software_list = download_data.selected_soft_list
        self._selected_soft_ver_list = download_data.selected_soft_ver_list
        self._base_dir = download_data.base_dir
        self._software_mgr = download_data.software_mgr
        self._images_list = [soft_ver for soft_ver in self._selected_soft_ver_list if
                                    self._is_images(soft_ver)]
        self.warning_message = set()
        self._docker_available = False
        # check docker availability at init
        try:
            self._docker_available = self._check_docker_available()
        except Exception:
            # _check_docker_available already logs
            self._docker_available = False

    @staticmethod
    def _is_images(soft_ver: SoftwareVersion):
        return any(soft in soft_ver.name for soft in DockerImagesDownloader._IMAGES)

    def _check_docker_available(self) -> bool:
        """Check whether docker is available on PATH and callable.

        Returns True if docker exists and responds to `docker --version`, False otherwise.
        Logs meaningful messages consistent with project logging.
        """
        if not shutil.which('docker'):
            LOG.error('docker executable not found in PATH. Skipping docker image pull/save.')
            self.warning_message.add('[ASCEND][WARNING]: docker not found in PATH, docker images will be skipped.')
            return False
        try:
            proc = subprocess.run(['docker', '--version'], stdout=subprocess.PIPE, stderr=subprocess.PIPE, check=False)
            if proc.returncode != 0:
                LOG.error('docker --version failed: %s', proc.stderr.decode(errors='ignore'))
                self.warning_message.add('[ASCEND][WARNING]: docker not responding correctly, docker images will be skipped.')
                return False
            LOG.info('docker available: %s', proc.stdout.decode(errors='ignore').strip())
            return True
        except Exception as ex:
            LOG.error('Failed to execute docker --version: %s', str(ex))
            self.warning_message.add('[ASCEND][WARNING]: failed to execute docker, docker images will be skipped.')
            return False

    def _pull_image(self, image_tag: str, image_name: Optional[str]=None, arch=None) -> bool:
        """Pull docker image. On Windows, add platform/arch parameter.

        :param image_tag: image reference (e.g. repo/name:tag)
        :param image_name: optional filename to save as—only used for logging
        :param arch: architecture string like 'x86_64' or 'aarch64'
        :return: True when pull succeeded
        """
        if not self._docker_available:
            LOG.error('docker not available, cannot pull image %s', image_tag)
            raise RuntimeError('docker not available, cannot pull image')

        cmd = ['docker', 'pull']
        # when running on Windows, docker may need explicit platform mapping
        if platform.system().lower() == 'windows':
            plat = None
            if arch == 'aarch64':
                plat = 'linux/arm64/v8'
            elif arch == 'x86_64':
                plat = 'linux/amd64'
            if plat:
                cmd.extend(['--platform', plat])

        cmd.append(image_tag)
        LOG.info('pull image cmd: %s', ' '.join(cmd))
        output_lines = []
        try:
            # Stream docker pull output so user sees progress
            proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
            # Read and display lines as they come
            for line in proc.stdout:
                if line is None:
                    continue
                line = line.rstrip()
                if line:
                    output_lines.append(line)
                    print(line)
                    LOG.info('docker pull: %s', line)
            proc.wait()
            if proc.returncode != 0:
                output_str = '\n'.join(output_lines)
                raise RuntimeError(f'Failed to pull image {image_tag} with exit code {proc.returncode}:\n{output_str}')
            LOG.info('Pulled image %s successfully', image_tag)
            return True
        except Exception as ex:
            LOG.error('Exception when pulling image %s: %s', image_tag, str(ex))
            raise

    def _save_image(self, image_tag: str, image_name: Optional[str]=None, download_dir: str = None, arch=None) -> bool:
        """Save docker image to a tar file named `image_name` under `download_dir`.

        :param image_tag: image reference
        :param image_name: target file name
        :param download_dir: target directory
        :return: True on success
        """
        if not self._docker_available:
            LOG.error('docker not available, cannot save image %s', image_tag)
            raise RuntimeError('docker not available, cannot save image')

        if not download_dir:
            download_dir = os.getcwd()
        if not os.path.exists(download_dir):
            os.makedirs(download_dir, mode=0o750, exist_ok=True)

        if not image_name:
            # fallback filename from tag
            image_name = image_tag.replace('/', '_').replace(':', '_') + '.tar'
        
        image_name = image_name + '_' + (arch if arch else '')

        if not image_name.endswith('.tar.gz'):
            image_name += '.tar.gz'
        dst_path = os.path.abspath(os.path.join(download_dir, image_name))

        # System-specific command construction
        system = platform.system().lower()
        if system == 'linux':
            # Check if gzip command exists in Linux environment
            try:
                # Verify gzip is available
                subprocess.run(['which', 'gzip'], check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
            except subprocess.CalledProcessError:
                raise RuntimeError('gzip command not found in Linux environment, cannot compress image file')
            
            # Linux: docker save output to gzip for direct compression (no intermediate files)
            cmd = [
                'docker', 'save', image_tag,
                '|', 'gzip', '-f'  # -f force overwrite existing file
            ]
            # Concatenate output path (shell pipe redirection)
            cmd_str = ' '.join(cmd) + f' > "{dst_path}"'
            shell = True  # Linux needs shell support for pipe
        elif system == 'windows':
            # Windows: PowerShell command (native Compress-Archive, no additional dependencies)
            # Key: docker save output to pipe, compress to tar.gz via Compress-Archive
            # Note: Compress-Archive requires absolute path for -DestinationPath to avoid PowerShell path issues
            docker_cmd = ['docker', 'save', image_tag]
            # Append Windows architecture parameters (e.g., --platform linux/amd64)
            if arch:
                plat_map = {'aarch64': 'linux/arm64/v8', 'x86_64': 'linux/amd64'}
                plat = plat_map.get(arch)
                if plat:
                    docker_cmd.insert(2, '--platform')
                    docker_cmd.insert(3, plat)
            
            # Build PowerShell command string (escape quotes to ensure correct path)
            docker_cmd_str = ' '.join([f'"{arg}"' for arg in docker_cmd])  # Add quotes to each parameter (handle spaces)
            powershell_cmd = (
                f'& {docker_cmd_str} | Compress-Archive -DestinationPath "{dst_path}" -Force -CompressionLevel Optimal'
            )
            # Windows must execute via powershell.exe with pipe enabled
            cmd = ['powershell.exe', '-Command', powershell_cmd]
            cmd_str = ' '.join(cmd)
            shell = False  # No need for shell=True when using powershell.exe on Windows
        else:
            raise RuntimeError(f'Unsupported operating system: {platform.system()}, only Linux/Windows are supported')

        # Execute command and capture logs
        LOG.info('Start saving image to tar.gz: %s', cmd_str)
        try:
            # Execute command (stream output logs to avoid blocking)
            proc = subprocess.Popen(
                cmd if system == 'windows' else cmd_str,
                stdout=subprocess.PIPE,
                stderr=subprocess.STDOUT,
                text=True,
                shell=shell,
                encoding='utf-8'  # Resolve Chinese log garbled issue on Windows
            )

            # Print and record logs in real-time
            for line in proc.stdout:
                if line.strip():
                    print(line.strip())
                    LOG.info('Execution log: %s', line.strip())
            
            # Wait for command completion and check return code
            proc.wait()
            if proc.returncode != 0:
                raise RuntimeError(f'Command execution failed (return code: {proc.returncode}), command: {cmd_str}')

            # Verify target file exists and has reasonable size
            if not os.path.exists(dst_path):
                raise RuntimeError(f'Command executed successfully but generated tar.gz file not found: {dst_path}')
            if os.path.getsize(dst_path) < 1024:  # Less than 1KB is considered abnormal (empty/corrupted file)
                raise RuntimeError(f'tar.gz file is too small (possibly corrupted): {dst_path}, size: {os.path.getsize(dst_path)} bytes')

            LOG.info('Image saved to tar.gz successfully: %s', dst_path)
            return True
        except Exception as ex:
            # Clean up potentially incomplete files on exception
            if os.path.exists(dst_path):
                os.remove(dst_path)
                LOG.warning('Cleaned up incomplete tar.gz file: %s', dst_path)
            LOG.error('Failed to save image to tar.gz: %s', str(ex))
            raise

    def _pull_and_save_image(self, soft_ver: SoftwareVersion, arch) -> List[DownloadFileInfo]:
        other_pkgs = self._software_mgr.get_software_other(soft_ver.name, soft_ver.version)
        for item in other_pkgs:
            print(f"analysis results: filename: {item.filename}")
            # _pull_image/_save_image will raise on failure to ensure the entire flow stops
            try:
                download_dir = os.path.join(self._base_dir, item.dest) if item.dest else os.path.join(self._base_dir, "resources", soft_ver.name)
                self._pull_image(item.url, item.filename, arch)
                self._save_image(item.url, item.filename, download_dir, arch)
                LOG.info(f'{item.filename} pulled and saved with the tag {item.url}')
            except Exception as ex:
                LOG.error(f'Failed to pull and save image {item.url} as {item.filename}: {str(ex)}')
                raise DownloadError(item.url, download_dir, str(ex))
        return 

    def collect_docker_images(self):

        # If no images need processing, do nothing (avoid docker checks)
        if not self._images_list:
            return 

        # If images exist but docker is not available, raise to abort downloads
        if not self._docker_available:
            msg = 'docker not available. If you are using Windows, please install Docker Desktop.'
            LOG.error(msg)
            # Record as an error so caller can inspect messages if needed
            self.warning_message.add('[ASCEND][ERROR]: docker not found or not responding.')
            # Raise an exception to stop further download tasks
            raise RuntimeError(msg)

        for soft_ver in self._images_list:
            self._pull_and_save_image(soft_ver, self._arch)




    

    

    
 

 